import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

tf.set_random_seed(777)  # for reproducibility

x = np.linspace(-1, 1, 300)
y = x * x - 0.5 + np.random.normal(0, 0.05, 300)

x2 = x.reshape([300, 1])
y2 = y.reshape([300, 1])

plt.plot(x, y)
plt.show()

X = tf.placeholder(tf.float32, shape=[None, 1])
Y = tf.placeholder(tf.float32, shape=[None, 1])

W1 = tf.Variable(tf.random_normal([1, 20]), name='weight1')
b1 = tf.Variable(tf.random_normal([20]), name='bias1')
layer1 = tf.sigmoid(tf.matmul(X, W1) + b1)
# layer1 = tf.nn.relu(tf.matmul(X, W1) + b1)

W2 = tf.Variable(tf.random_normal([20, 1]), name='weight2')
b2 = tf.Variable(tf.random_normal([1]), name='bias2')
hypothesis = tf.matmul(layer1, W2) + b2

cost = tf.reduce_mean(tf.reduce_sum(tf.square(hypothesis - Y), 1))

train = tf.train.GradientDescentOptimizer(learning_rate=0.1).minimize(cost)

sess = tf.Session()
sess.run(tf.global_variables_initializer())

for step in range(1001):
    cost_val, hy_val, train_val = sess.run(
        [cost, hypothesis, train], feed_dict={X: x2, Y: y2})
    if step % 50 == 0:
        print(step, "Cost: ", cost_val)

print(y2[299:300],": ", hy_val[299:300])
